/*
 * @Author: lixu lixu@puchigames.com
 * @Date: 2025-09-18 11:20:00
 * @LastEditors: lixu lixu@puchigames.com
 * @LastEditTime: 2025-09-18 11:20:00
 * @FilePath: /go-helper/ixLog/trace.go
 * @Description: 分布式追踪功能实现，支持TraceID、SpanID、调用链追踪等
 */

package ixLog

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"
)

// TraceManager 追踪管理器
type TraceManager struct {
	config    *Config
	sampler   Sampler
	spanStore *SpanStore
	mutex     sync.RWMutex
}

// NewTraceManager 创建追踪管理器
func NewTraceManager(config *Config) *TraceManager {
	return &TraceManager{
		config:    config,
		sampler:   NewProbabilitySampler(config.SampleRate),
		spanStore: NewSpanStore(1000), // 默认存储1000个span
	}
}

// StartTrace 开始追踪
func (tm *TraceManager) StartTrace(ctx context.Context, operationName string) (context.Context, *Span) {
	traceID := generateTraceID()
	spanID := generateSpanID()

	span := &Span{
		TraceID:       traceID,
		SpanID:        spanID,
		OperationName: operationName,
		StartTime:     time.Now(),
		Tags:          make(map[string]interface{}),
		Logs:          make([]*LogEntry, 0),
	}

	// 检查是否采样
	if tm.sampler.ShouldSample(traceID) {
		span.Sampled = true
		tm.spanStore.AddSpan(span)
	}

	// 将span添加到上下文
	traceCtx := &TraceContext{
		TraceID: traceID,
		SpanID:  spanID,
		Baggage: make(map[string]interface{}),
	}

	ctx = context.WithValue(ctx, TraceContextKey, traceCtx)
	return ctx, span
}

// StartChildSpan 开始子追踪
func (tm *TraceManager) StartChildSpan(ctx context.Context, operationName string) (context.Context, *Span) {
	parentCtx := GetTraceContext(ctx)
	if parentCtx == nil {
		return tm.StartTrace(ctx, operationName)
	}

	spanID := generateSpanID()

	span := &Span{
		TraceID:       parentCtx.TraceID,
		SpanID:        spanID,
		ParentSpanID:  parentCtx.SpanID,
		OperationName: operationName,
		StartTime:     time.Now(),
		Tags:          make(map[string]interface{}),
		Logs:          make([]*LogEntry, 0),
	}

	// 检查是否采样
	if tm.sampler.ShouldSample(parentCtx.TraceID) {
		span.Sampled = true
		tm.spanStore.AddSpan(span)
	}

	// 更新上下文
	newTraceCtx := &TraceContext{
		TraceID:      parentCtx.TraceID,
		SpanID:       spanID,
		ParentSpanID: parentCtx.SpanID,
		Baggage:      parentCtx.Baggage,
	}

	ctx = context.WithValue(ctx, TraceContextKey, newTraceCtx)
	return ctx, span
}

// FinishSpan 结束追踪
func (tm *TraceManager) FinishSpan(span *Span) {
	if span == nil {
		return
	}

	span.EndTime = time.Now()
	span.Duration = span.EndTime.Sub(span.StartTime)

	if span.Sampled {
		tm.spanStore.UpdateSpan(span)
	}
}

// GetTracesByID 根据TraceID获取所有span
func (tm *TraceManager) GetTracesByID(traceID string) []*Span {
	return tm.spanStore.GetSpansByTraceID(traceID)
}

// Span 追踪跨度
type Span struct {
	TraceID       string                 `json:"trace_id"`
	SpanID        string                 `json:"span_id"`
	ParentSpanID  string                 `json:"parent_span_id,omitempty"`
	OperationName string                 `json:"operation_name"`
	StartTime     time.Time              `json:"start_time"`
	EndTime       time.Time              `json:"end_time,omitempty"`
	Duration      time.Duration          `json:"duration,omitempty"`
	Tags          map[string]interface{} `json:"tags,omitempty"`
	Logs          []*LogEntry            `json:"logs,omitempty"`
	Sampled       bool                   `json:"sampled"`
	Error         error                  `json:"error,omitempty"`
}

// SetTag 设置标签
func (s *Span) SetTag(key string, value interface{}) *Span {
	if s.Tags == nil {
		s.Tags = make(map[string]interface{})
	}
	s.Tags[key] = value
	return s
}

// LogFields 记录字段
func (s *Span) LogFields(fields map[string]interface{}) {
	entry := &LogEntry{
		Timestamp: time.Now(),
		Level:     TRACE,
		Message:   "span log",
		Fields:    fields,
		TraceID:   s.TraceID,
		SpanID:    s.SpanID,
	}
	s.Logs = append(s.Logs, entry)
}

// SetError 设置错误
func (s *Span) SetError(err error) *Span {
	s.Error = err
	s.SetTag("error", true)
	s.SetTag("error.message", err.Error())
	return s
}

// Sampler 采样器接口
type Sampler interface {
	ShouldSample(traceID string) bool
}

// ProbabilitySampler 概率采样器
type ProbabilitySampler struct {
	rate float64
}

// NewProbabilitySampler 创建概率采样器
func NewProbabilitySampler(rate float64) *ProbabilitySampler {
	if rate < 0 {
		rate = 0
	}
	if rate > 1 {
		rate = 1
	}
	return &ProbabilitySampler{rate: rate}
}

// ShouldSample 判断是否应该采样
func (s *ProbabilitySampler) ShouldSample(traceID string) bool {
	if s.rate >= 1.0 {
		return true
	}
	if s.rate <= 0.0 {
		return false
	}

	// 基于traceID的哈希值进行采样决策
	hash := hashString(traceID)
	threshold := uint64(s.rate * float64(^uint64(0)))
	return hash < threshold
}

// SpanStore 跨度存储
type SpanStore struct {
	spans   map[string]*Span
	maxSize int
	mutex   sync.RWMutex
	cleanup chan string
}

// NewSpanStore 创建跨度存储
func NewSpanStore(maxSize int) *SpanStore {
	store := &SpanStore{
		spans:   make(map[string]*Span),
		maxSize: maxSize,
		cleanup: make(chan string, 100),
	}

	// 启动清理协程
	go store.cleanupLoop()

	return store
}

// AddSpan 添加跨度
func (ss *SpanStore) AddSpan(span *Span) {
	ss.mutex.Lock()
	defer ss.mutex.Unlock()

	// 检查容量
	if len(ss.spans) >= ss.maxSize {
		// 删除最旧的span
		ss.removeOldestSpan()
	}

	ss.spans[span.SpanID] = span
}

// UpdateSpan 更新跨度
func (ss *SpanStore) UpdateSpan(span *Span) {
	ss.mutex.Lock()
	defer ss.mutex.Unlock()

	ss.spans[span.SpanID] = span

	// 如果span已结束，标记为可清理
	if !span.EndTime.IsZero() {
		select {
		case ss.cleanup <- span.SpanID:
		default:
		}
	}
}

// GetSpan 获取跨度
func (ss *SpanStore) GetSpan(spanID string) *Span {
	ss.mutex.RLock()
	defer ss.mutex.RUnlock()

	return ss.spans[spanID]
}

// GetSpansByTraceID 根据TraceID获取所有跨度
func (ss *SpanStore) GetSpansByTraceID(traceID string) []*Span {
	ss.mutex.RLock()
	defer ss.mutex.RUnlock()

	var spans []*Span
	for _, span := range ss.spans {
		if span.TraceID == traceID {
			spans = append(spans, span)
		}
	}

	return spans
}

// removeOldestSpan 移除最旧的跨度
func (ss *SpanStore) removeOldestSpan() {
	var oldestSpanID string
	var oldestTime time.Time

	for spanID, span := range ss.spans {
		if oldestSpanID == "" || span.StartTime.Before(oldestTime) {
			oldestSpanID = spanID
			oldestTime = span.StartTime
		}
	}

	if oldestSpanID != "" {
		delete(ss.spans, oldestSpanID)
	}
}

// cleanupLoop 清理循环
func (ss *SpanStore) cleanupLoop() {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case spanID := <-ss.cleanup:
			// 延迟删除已完成的span
			time.AfterFunc(5*time.Minute, func() {
				ss.mutex.Lock()
				delete(ss.spans, spanID)
				ss.mutex.Unlock()
			})
		case <-ticker.C:
			// 定期清理超时的span
			ss.cleanupExpiredSpans()
		}
	}
}

// cleanupExpiredSpans 清理过期的跨度
func (ss *SpanStore) cleanupExpiredSpans() {
	ss.mutex.Lock()
	defer ss.mutex.Unlock()

	cutoff := time.Now().Add(-10 * time.Minute) // 10分钟前的span被认为过期

	for spanID, span := range ss.spans {
		if span.StartTime.Before(cutoff) {
			delete(ss.spans, spanID)
		}
	}
}

// GetTraceContext 从上下文获取追踪信息
func GetTraceContext(ctx context.Context) *TraceContext {
	if ctx == nil {
		return nil
	}

	if traceCtx, ok := ctx.Value(TraceContextKey).(*TraceContext); ok {
		return traceCtx
	}

	return nil
}

// WithTraceContext 将追踪信息添加到上下文
func WithTraceContext(ctx context.Context, traceCtx *TraceContext) context.Context {
	return context.WithValue(ctx, TraceContextKey, traceCtx)
}

// GetCallStack 获取调用栈
func GetCallStack(skip, depth int) []string {
	if depth <= 0 {
		return nil
	}

	var stack []string
	for i := skip; i < skip+depth; i++ {
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}

		funcName := "unknown"
		if fn := runtime.FuncForPC(pc); fn != nil {
			funcName = fn.Name()
		}

		// 简化文件路径
		if idx := strings.LastIndex(file, "/"); idx != -1 {
			file = file[idx+1:]
		}

		stack = append(stack, fmt.Sprintf("%s:%d %s", file, line, funcName))
	}

	return stack
}

// GetCallerInfo 获取调用者信息
func GetCallerInfo(skip int) *CallerInfo {
	pc, file, line, ok := runtime.Caller(skip)
	if !ok {
		return nil
	}

	funcName := "unknown"
	if fn := runtime.FuncForPC(pc); fn != nil {
		funcName = fn.Name()
	}

	return &CallerInfo{
		File:     file,
		Line:     line,
		Function: funcName,
	}
}

// generateTraceID 生成追踪ID
func generateTraceID() string {
	bytes := make([]byte, 16)
	rand.Read(bytes)
	return hex.EncodeToString(bytes)
}

// generateSpanID 生成跨度ID
func generateSpanID() string {
	bytes := make([]byte, 8)
	rand.Read(bytes)
	return hex.EncodeToString(bytes)
}

// hashString 计算字符串哈希值
func hashString(s string) uint64 {
	hash := uint64(0)
	for _, char := range s {
		hash = hash*31 + uint64(char)
	}
	return hash
}

// InjectTraceHeaders 注入追踪头信息
func InjectTraceHeaders(ctx context.Context, headers map[string]string) {
	traceCtx := GetTraceContext(ctx)
	if traceCtx == nil {
		return
	}

	headers["X-Trace-ID"] = traceCtx.TraceID
	headers["X-Span-ID"] = traceCtx.SpanID
	if traceCtx.ParentSpanID != "" {
		headers["X-Parent-Span-ID"] = traceCtx.ParentSpanID
	}

	// 注入Baggage
	if len(traceCtx.Baggage) > 0 {
		var baggage []string
		for key, value := range traceCtx.Baggage {
			baggage = append(baggage, fmt.Sprintf("%s=%v", key, value))
		}
		headers["X-Trace-Baggage"] = strings.Join(baggage, ",")
	}
}

// ExtractTraceHeaders 提取追踪头信息
func ExtractTraceHeaders(headers map[string]string) *TraceContext {
	traceID := headers["X-Trace-ID"]
	if traceID == "" {
		return nil
	}

	traceCtx := &TraceContext{
		TraceID:      traceID,
		SpanID:       headers["X-Span-ID"],
		ParentSpanID: headers["X-Parent-Span-ID"],
		Baggage:      make(map[string]interface{}),
	}

	// 解析Baggage
	if baggage := headers["X-Trace-Baggage"]; baggage != "" {
		pairs := strings.Split(baggage, ",")
		for _, pair := range pairs {
			parts := strings.SplitN(pair, "=", 2)
			if len(parts) == 2 {
				traceCtx.Baggage[parts[0]] = parts[1]
			}
		}
	}

	return traceCtx
}

// SetBaggage 设置行李数据
func SetBaggage(ctx context.Context, key string, value interface{}) context.Context {
	traceCtx := GetTraceContext(ctx)
	if traceCtx == nil {
		traceCtx = &TraceContext{
			Baggage: make(map[string]interface{}),
		}
	} else {
		// 复制上下文避免修改原始数据
		newTraceCtx := *traceCtx
		newTraceCtx.Baggage = make(map[string]interface{})
		for k, v := range traceCtx.Baggage {
			newTraceCtx.Baggage[k] = v
		}
		traceCtx = &newTraceCtx
	}

	traceCtx.Baggage[key] = value
	return context.WithValue(ctx, TraceContextKey, traceCtx)
}

// GetBaggage 获取行李数据
func GetBaggage(ctx context.Context, key string) interface{} {
	traceCtx := GetTraceContext(ctx)
	if traceCtx == nil {
		return nil
	}

	return traceCtx.Baggage[key]
}

// GetAllBaggage 获取所有行李数据
func GetAllBaggage(ctx context.Context) map[string]interface{} {
	traceCtx := GetTraceContext(ctx)
	if traceCtx == nil {
		return nil
	}

	// 返回副本避免外部修改
	baggage := make(map[string]interface{})
	for k, v := range traceCtx.Baggage {
		baggage[k] = v
	}

	return baggage
}

// TraceMiddleware HTTP追踪中间件
func TraceMiddleware(tm *TraceManager) func(next func()) func() {
	return func(next func()) func() {
		return func() {
			// 这里可以实现HTTP中间件逻辑
			// 由于这是一个通用的日志包，这里提供一个简单的示例
			ctx := context.Background()
			ctx, span := tm.StartTrace(ctx, "http_request")
			defer tm.FinishSpan(span)

			span.SetTag("component", "http")
			span.SetTag("http.method", "GET") // 示例

			next()
		}
	}
}

// WithTraceLogger 为上下文添加带追踪的日志器
func WithTraceLogger(ctx context.Context, logger Logger) context.Context {
	traceCtx := GetTraceContext(ctx)
	if traceCtx != nil {
		logger = logger.WithTrace(traceCtx.TraceID, traceCtx.SpanID, traceCtx.ParentSpanID)
	}

	return context.WithValue(ctx, LoggerContextKey, logger)
}

// GetTraceLogger 从上下文获取带追踪的日志器
func GetTraceLogger(ctx context.Context) Logger {
	if logger, ok := ctx.Value(LoggerContextKey).(Logger); ok {
		return logger
	}
	return nil
}
